/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.common;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicLong;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.mapreduce.SssException;

public class MultipleBufferedQueue<INPUT extends RoughSizeGetable, OUTPUT>
  implements DataReceiver<INPUT> {
  private static Logger logger = LoggerFactory.getLogger(MultipleBufferedQueue.class);

  private final ActiveBufferIndex active = new ActiveBufferIndex();
  private final int multiplicity;
  private final List<Buffer<INPUT, OUTPUT>> bufferList;
  private final Executor executor;
  private final DataReceiver<OUTPUT> merger;
  private final long limitNBytes;
  private final AtomicLong inputCount  = new AtomicLong(0);
  private final AtomicLong outputCount = new AtomicLong(0);

  public static <INPUT extends RoughSizeGetable> MultipleBufferedQueue<INPUT, List<INPUT>>
  createList(int multiplicity,
             long limitNBytes,
             Executor executor,
             DataReceiver<List<INPUT>> merger) {
    MultipleBufferedQueue<INPUT, List<INPUT>> bq =
      new MultipleBufferedQueue<INPUT, List<INPUT>>(multiplicity, limitNBytes, executor, merger);

    for (int i = 0; i < multiplicity; i++) {
      bq.bufferList.add(new ListBuffer<INPUT>(bq));
    }
    return bq;
  }

  private MultipleBufferedQueue(int multiplicity,
                                long limitNBytes,
                                Executor executor,
                                DataReceiver<OUTPUT> merger) {
    assert multiplicity > 0: "multiplicity should be positive integer.";

    this.multiplicity = multiplicity;
    this.limitNBytes = limitNBytes;
    this.executor = executor;
    this.merger = new RefCountDataReceiver<OUTPUT>(merger, multiplicity);
    this.bufferList = new ArrayList<Buffer<INPUT, OUTPUT>>(multiplicity);
  }

  @Override
  public void put(INPUT bucket) throws SssException {
    checkAlreadyError();
    try {
      inputCount.incrementAndGet();
      synchronized (active) {
        Buffer<INPUT, OUTPUT> buffer = bufferList.get(active.index);
        if (buffer.put(bucket)) {;
          active.next();
        }
      }
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public void close() throws SssException {
    logger.debug("close");
    checkAlreadyError();
    synchronized (active) {
      for (Buffer<INPUT, OUTPUT> buffer: bufferList) {
        buffer.flush(true);
      }
    }
  }

  @Override
  public void setError(Throwable e) {
    merger.setError(e);
  }

  @Override
  public Throwable getError() {
    return merger.getError();
  }

  private void checkAlreadyError() throws ErrorDataReceiverException {
    Throwable e = merger.getError();
    if (e != null) {
      throw new ErrorDataReceiverException(e);
    }
  }

  private class ActiveBufferIndex {
    private int index;
    private ActiveBufferIndex() {
      this.index = 0;
    }
    private void next() {
      index = (index + 1) % multiplicity;
    }
  }

  private abstract static class Buffer<INPUT extends RoughSizeGetable, OUTPUT> {
    private final MultipleBufferedQueue<INPUT, OUTPUT> owner;
    private boolean flushing = false;
    private boolean last = false;
    private long nbytes = 0L;

    public Buffer(MultipleBufferedQueue<INPUT, OUTPUT> owner) {
      this.owner = owner;
    }

    protected abstract void add(INPUT bucket);
    protected abstract OUTPUT get();
    protected abstract int size();
    protected abstract void clear();

    protected boolean isEmpty() {
      return size() == 0;
    }

    public void waitFlush() throws InterruptedException {
      while (flushing) {
        wait();
      }
    }

    /**
     *.
     * If buffer is full, this flushes buffer and returns true.
     */
    public synchronized boolean put(INPUT bucket) throws InterruptedException {
      waitFlush();
      add(bucket);
      nbytes += bucket.getRoughSize();
      if (nbytes >= owner.limitNBytes) {
        flush(false);
        return true;
      }
      return false;
    }

    /**
     * .
     * This is asynchronized.
     */
    public synchronized void flush(boolean last) {
      assert !this.last: "flush(true) is called more than once";
      this.last = last;
      if (flushing) {
        return;
      }
      if (isEmpty() && !last) {
        // this buffer is clean
        return;
      }

      logger.debug("queue size = {} items, {} bytes.", size(), nbytes);

      flushing = true;

      owner.executor.execute(new Runnable() {
        @Override
        public void run() {
          boolean completed = false;
          try {
            try {
              owner.outputCount.addAndGet(size());
              owner.merger.put(get());
            } finally {
              completed = flushCompleted();
            }
            if (completed) {
              owner.merger.close();
              logger.debug("closed input {} output {}",
                owner.inputCount.get(), owner.outputCount.get());
            }
          } catch (ErrorDataReceiverException e) {
            // Ignore
          } catch (Throwable e) {
            owner.merger.setError(e);
          }
        }
      });
    }

    public synchronized boolean flushCompleted() {
      assert flushing: "flushCompleted is called while flushing is false.";
      clear();
      nbytes = 0L;
      flushing = false;
      this.notifyAll();

      return last;
    }
  }

  private static class ListBuffer<INPUT extends RoughSizeGetable> extends Buffer<INPUT, List<INPUT>> {
    private final List<INPUT> buffer = new ArrayList<INPUT>();

    public ListBuffer(MultipleBufferedQueue<INPUT, List<INPUT>> owner) {
      super(owner);
    }

    @Override
    protected void add(INPUT bucket) {
      buffer.add(bucket);
    }

    @Override
    protected List<INPUT> get() {
      return buffer;
    }

    @Override
    protected int size() {
      return buffer.size();
    }

    @Override
    protected void clear() {
      buffer.clear();
    }
  }
}
